#include "ballot_generator.h"

ballot_generator::ballot_generator(){}

ballot_generator::~ballot_generator(){}

ballot_generator::ballot_generator(QString Dir){
    obj_Dir.setPath(Dir);
    obj_Dir.mkdir("BG");
    obj_Dir.cd("BG");

    create_key();
    pvid = new PVID_Authority(Dir);

    // Crea la base de datos que hace la funcion de BallotList(BL)
    if(!QFile::exists(obj_Dir.absoluteFilePath("BallotList").append(".db"))){
        QString cmd("sqlite3 -line ");
        cmd.append(obj_Dir.absoluteFilePath("BallotList")).append(".db ");
        cmd.append("'create table BallotList(PVID BLOB, M3 BLOB, B BLOB, es BLOB);'");
        system(cmd.toAscii().data());
    }

    // Crea la base de datos que hace la funcion de BallotGeneratorBulletinBoard(BGBB)
    if(!QFile::exists(obj_Dir.absoluteFilePath("BGBB").append(".db"))){
        QString cmd("sqlite3 -line ");
        cmd.append(obj_Dir.absoluteFilePath("BGBB")).append(".db ");
        cmd.append("'create table BGBB(H BLOB);'");
        system(cmd.toAscii().data());
    }

    if(!QSqlDatabase::contains("BLconex"))
        QSqlDatabase::addDatabase("QSQLITE","BLconex");
    if(!QSqlDatabase::contains("BGBBconex"))
        QSqlDatabase::addDatabase("QSQLITE","BGBBconex");
    if(!QSqlDatabase::contains("PVIDconex"))
        QSqlDatabase::addDatabase("QSQLITE","PVIDconex");
    if(!QSqlDatabase::contains("SubBLconex"))
        QSqlDatabase::addDatabase("QSQLITE","SubBLconex");

}

ballot_generator::ballot_generator(QString Dir, bool no_key){
    obj_Dir.setPath(Dir.append("/BG"));
    pvid = new PVID_Authority();
    if(!QSqlDatabase::contains("BLconex"))
        QSqlDatabase::addDatabase("QSQLITE","BLconex");
    if(!QSqlDatabase::contains("BGBBconex"))
        QSqlDatabase::addDatabase("QSQLITE","BGBBconex");
    if(!QSqlDatabase::contains("PVIDconex"))
        QSqlDatabase::addDatabase("QSQLITE","PVIDconex");
    if(!QSqlDatabase::contains("SubBLconex"))
        QSqlDatabase::addDatabase("QSQLITE","SubBLconex");
}

void ballot_generator::SubBallotList(){
// Crea la conexion
    QSqlDatabase BL = QSqlDatabase::database("BLconex",false);
    BL.setDatabaseName(obj_Dir.absoluteFilePath("BallotList.db"));
    if(!BL.open()){
        return;
    }

// Crea el modelo
    QSqlTableModel BLmodel(0,BL);
    BLmodel.setTable("BallotList");
    BLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    BLmodel.select();

// Crea la base de datos que hace la funcion de SubBallotList
    if(QFile::exists(obj_Dir.absoluteFilePath("SubBallotList").append(".db"))){
        QFile::remove(obj_Dir.absoluteFilePath("SubBallotList.db"));
    }
    QString cmd("sqlite3 -line ");
    cmd.append(obj_Dir.absoluteFilePath("SubBallotList.db "));
    cmd.append("'create table SubBallotList(PVID BLOB, B BLOB);'");
    system(cmd.toAscii().data());


    QSqlDatabase SBL = QSqlDatabase::database("SubBLconex",false);
    SBL.setDatabaseName(obj_Dir.absoluteFilePath("SubBallotList.db"));
    if(!SBL.open()){
        return;
    }

// Crea el modelo
    QSqlTableModel SBLmodel(0,SBL);
    SBLmodel.setTable("SubBallotList");
    SBLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    SBLmodel.select();

//
    for(int i = 0;i<BLmodel.rowCount();i++){
        QString buffer(BLmodel.record(i).value("B").toString());
        QString B;
        shift_String(buffer);
        shift_String(buffer);
        SBLmodel.insertRows(0, 1);
        SBLmodel.setData(SBLmodel.index(0,SBLmodel.fieldIndex("PVID")),BLmodel.record(i).value("PVID"));
        shift_String(buffer);
        shift_String(buffer);
        for(int i=1;!buffer.isEmpty();i++){
            B.append(QString::number(i).append(". ")).append(shift_String(buffer)).append("&&");
        }
        SBLmodel.setData(SBLmodel.index(0,SBLmodel.fieldIndex("B")),B);
        SBLmodel.submitAll();
    }
}

bool ballot_generator::M1(QString &m1, QString &m2){
/*  Desencripta m1
    Divide el mensaje recibido en PVID, Ek y es
    Verifica el PVID para comprobar que es un votante válido
    Guarda la llave publica del votante (es)
    Firma
        -- abre la llave publica de BG
        -- une Ek con la llave publica de BG
        -- Firma la union anterior
    Encripta la firma anterior con la llave publica del KG

 */
    int PVID_size,Ek_size;
    QString Eb,ev,Ek;

// Desencripta m1
    if(!decryptMsg(m1,Eb)){
        return false;
    }

// Divide el contenido de Eb en cada una de sus partes
    PVID_size=shift_String(Eb).toInt();
    Ek_size=shift_String(Eb).toInt();

    PVID=shift_String(Eb,PVID_size);
    Ek=shift_String(Eb,Ek_size);
    es=Eb;

// Verifica la identidad del votante
    if(!pvid->PVID_verification(PVID,"PVID1")||pvid->isVoted(PVID)){
        return false;
    }

// Guarda la llave pública del votante
    Hex2Key(es,obj_Dir.absoluteFilePath("es"));

// Firma
    QString buffer,Db,eb;

    Key2Hex(obj_Dir.absoluteFilePath("eb"),eb);
    buffer = QString::number(Ek.size());
    buffer.append("&&");
    buffer.append(Ek);
    buffer.append(eb);

    Db = signMsg(buffer);

// Encripta M2
    m2=encryptMsg(Db,"ek");

    return true;
}

bool ballot_generator::M4(QString &m4, QString &m5){
/*  Desencripta M4
    Verifica la firma del KG
        -- Divide la informacion en M3 y ek
        -- verifica que la ek recibida es igual a la ek almacenada
    Crea una boleta dinámica (B)
    Publica el hash H(B,es) en el BGBB
    Guardar PVID,M3,B y es en el BallotList
    Genera M5
        -- Une M3, B y eb
        -- Firma la unión anterior
        -- Encripta la firma con la llave publica recibida del votante*/

    // Desencripta m4
    QString Eb,Dk,m3,ek,hexKey,B,Db,Es;
    int obj_size;

    if(!decryptMsg(m4,Eb)){

        return false;
    }

    // Verifica la firma del Key Generator y divide la info en M3 y ek
    if(!verifyMsg(Eb,"ek",Dk)){
        return false;
    }

    //
    obj_size=shift_String(Dk).toInt();
    m3=shift_String(Dk,obj_size);
    ek=Dk;

    if(!Key2Hex("ek",hexKey)){
        return false;
    }
    if(ek!=hexKey){
        return false;
    }

    // Crea una Boleta dinamica
    generate_ballot(B);

    // Guarda el hashes H(B,es)
    if(!Save_Hash(B,"es")){
        return false;
    }

    //Guardar PVID,M3,B y es en el BallotList
    if(!Save_in_BL(PVID,m3,B,es)){
        return false;
    }

    // Genera M5
    if(!Key2Hex("eb",hexKey)){
        return false;
    }

    Db=QString::number(m3.size());
    Db.append("&&");
    Db.append(QString::number(B.size())).append("&&");
    Db.append(m3);
    Db.append(B);
    Db.append(hexKey);

    Es = signMsg(Db);

    m5 = encryptMsg(Es,"es");
    return true;
}

QString ballot_generator::encryptMsg(QString Data,QString keyName){
    QString out;
    string plainMsg,encryptedMsg,buf;

//Convierto de QString a std::string
//En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y despues no se puede desencriptar
    plainMsg=Data.toStdString();

    RSA::PublicKey pub;

    LoadPublicKey(keyName,pub);

    RSAES_OAEP_SHA_Encryptor e( pub);

    while(plainMsg.size()>0){
        buf.clear();
        if(plainMsg.size()<=86){
            buf=plainMsg;
            plainMsg.clear();

        }
        else{
            for(int j=0;j<86; j++){
                buf+=plainMsg.at(j);
            }
            plainMsg.erase(0,86);
        }

        StringSource(buf, true,
            new PK_EncryptorFilter( rng, e,
                new StringSink( encryptedMsg )
            ) // PK_EncryptorFilter
         ); // StringSource
    }

    return QString::fromStdString(encryptedMsg);
}

bool ballot_generator::decryptMsg(QString &Data,QString &decryptedMsg){
    string decrypted_Msg,buf;
    RSA::PrivateKey pri;

    LoadPrivateKey("db",pri);

    RSAES_OAEP_SHA_Decryptor d( pri );

    //Convierto de QString a std::string
    //En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y no se puede desencriptar
    try{
        while(Data.size()>0){
            buf=QString(shift_String(Data,128)).toStdString();

            StringSource(buf, true,
                new PK_DecryptorFilter( rng, d,
                    new StringSink( decrypted_Msg )
                ) // PK_DecryptorFilter
             ); // StringSource

        }
    }
    catch(...){
        return false;
    }
    decryptedMsg = QString::fromStdString(decrypted_Msg);
    return true;
}

QString ballot_generator::signMsg(QString &msg){
    string plainMsg,signedMsg;

//Convierto de QString a std::string
//En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y despues no se puede desencriptar
    plainMsg=msg.toStdString();

// Carga la llave privada del objeto
    RSA::PrivateKey pri;
    LoadPrivateKey("db",pri);

    RSASS<PSSR, SHA>::Signer signer(pri);

    StringSource(plainMsg, true,
        new SignerFilter(rng, signer,
            new StringSink(signedMsg),
            true // putMessage for recovery
        ) // SignerFilter
     ); // StringSource*/

    return QString::fromStdString(signedMsg);
}

bool ballot_generator::verifyMsg(QString msg,QString keyToUse, QString &verifiedMsg){
    string signedMsg, verified_Msg;

    signedMsg=msg.toStdString();

    RSA::PublicKey pub;
    LoadPublicKey(keyToUse,pub);

//    CryptoPP::RSASSA_PKCS1v15_SHA_Verifier verifier(pub);
    RSASS<PSSR, SHA>::Verifier verifier(pub);

    try{
        StringSource(signedMsg, true,
            new SignatureVerificationFilter(
                verifier,
                new StringSink(verified_Msg),
                SignatureVerificationFilter::THROW_EXCEPTION | SignatureVerificationFilter::PUT_MESSAGE
           ) // SignatureVerificationFilter
        ); // StringSource
    } catch(Exception e){
        return false;
    }
    verifiedMsg=QString::fromStdString(verified_Msg);

return true;
}

QString ballot_generator::shift_String(QString &msg){
    QString buffer;

    int shift = msg.indexOf("&&");

    for(int j=0;j<shift; j++){
        buffer.append(msg.at(j));
    }

    msg.remove(0,shift);
    msg.remove(0,2); // borra &&
    return buffer;
}

QString ballot_generator::shift_String(QString &msg,int max){
    QString buffer;

    for(int j=0;j<max; j++){
        buffer.append(msg.at(j));
    }

    msg.remove(0,max);
    return buffer;
}

QString ballot_generator::Hash(QString B, QString keyname){
    // Abre el archivo keyname y carga su contenido en un QByteArray BA_key2
        QFile file(obj_Dir.absoluteFilePath(keyname));
        file.open(QIODevice::ReadOnly);
        QByteArray BA_key = file.readAll();
        file.close();

    // Genera la funcion hash de las llaves key1 y key2
        SHA sha;
        string output;

        sha.Update((byte const*)B.toAscii().data(), B.size());
        sha.Update((byte const*)BA_key.data(), BA_key.size());

        SecByteBlock sbbDigest(sha.DigestSize());

        sha.Final(sbbDigest.begin());
        HexEncoder enc(new StringSink(output));
        enc.Put(sbbDigest.begin(), sbbDigest.size());
        enc.MessageEnd();

        return QString::fromStdString(output);
}

bool ballot_generator::Hex2Key(QString keydata, QString keyname){
    QFile file(obj_Dir.absoluteFilePath(keyname));
    if(!file.open(QIODevice::WriteOnly)){
        return false;}

//    QByteArray BA = file.readAll();

    string decodedKey;

    HexDecoder deco(new StringSink(decodedKey));
    deco.Put((byte*)keydata.toStdString().data(),keydata.size());
    deco.MessageEnd();

    file.write(decodedKey.data(),decodedKey.size());

    file.close();
    return true;
}

bool ballot_generator::Key2Hex(QString keyname, QString &keyHex){
    QFile file(obj_Dir.absoluteFilePath(keyname));
    if(!file.open(QIODevice::ReadOnly)){
        return false;}

    QByteArray BA = file.readAll();

    string encodedKey;

    HexEncoder hex(new StringSink(encodedKey));
    hex.Put((const byte*)BA.data(),BA.size());
    hex.MessageEnd();

    keyHex=QString::fromStdString(encodedKey);

    return true;
}

bool ballot_generator::Save_Hash(QString B, QString VoterKey){
/*  Crear la conexion con la base de datos
    Crear el modelo que manipule la base de datos
    Crear el hash H(B,es)
    Copiar el hash en la base de datos
  */
// Crea la conexion
    QSqlDatabase BGBB = QSqlDatabase::database("BGBBconex",false);
    BGBB.setDatabaseName(obj_Dir.absoluteFilePath("BGBB.db"));

    if(!BGBB.open()){
        return false;
    }

// Crea el modelo
    QSqlTableModel BGBBmodel(0,BGBB);
    BGBBmodel.setTable("BGBB");
    BGBBmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    BGBBmodel.select();

// Crea el hash
    QString H = Hash(B,VoterKey);


// Copia las llaves en la base de datos
    int row = 0;
    BGBBmodel.insertRows(row, 1);
    BGBBmodel.setData(BGBBmodel.index(row, BGBBmodel.fieldIndex("H")),H);
    BGBBmodel.submitAll();
    return true;
}

bool ballot_generator::Save_in_BL(QString PVID, QString M3, QString B, QString es){
/*  Crear la conexion con la base de datos
    Crear el modelo que manipule la base de datos

    Copiar los datos en la base de datos
  */

// Crea la conexion
    QSqlDatabase BL = QSqlDatabase::database("BLconex",false);
    BL.setDatabaseName(obj_Dir.absoluteFilePath("BallotList.db"));

    if(!BL.open()){
        return false;
    }

// Crea el modelo
    QSqlTableModel BLmodel(0,BL);
    BLmodel.setTable("BallotList");
    BLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    BLmodel.select();

// Copia las llaves en la base de datos
    int row = 0;
    BLmodel.insertRows(row, 1);
    BLmodel.setData(BLmodel.index(row, BLmodel.fieldIndex("PVID")),PVID);
    BLmodel.setData(BLmodel.index(row, BLmodel.fieldIndex("M3")),M3);
    BLmodel.setData(BLmodel.index(row, BLmodel.fieldIndex("B")),B);
    BLmodel.setData(BLmodel.index(row, BLmodel.fieldIndex("es")),es);
    BLmodel.submitAll();

    return true;
}

bool ballot_generator::Search_in_BGBB(QString Hash){
    // Crea la conexion
    QSqlDatabase BGBB = QSqlDatabase::database("BGBBconex",false);
    BGBB.setDatabaseName(obj_Dir.absoluteFilePath("BGBB.db"));

    if(!BGBB.open()){
        return false;
    }

    // Crea el modelo
    QSqlTableModel BGBBmodel(0,BGBB);
    BGBBmodel.setTable("BGBB");
    BGBBmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    BGBBmodel.select();

    // verifica si Hash se encuentra en la base de datos
    for (int i = 0; i <= BGBBmodel.rowCount(); ++i) {
        if (Hash==BGBBmodel.record(i).value("H").toString()){
            return true;
        }
    }
    return false;
}

void ballot_generator::generate_ballot(QString &Ballot){
    int pointer;
    Ballot.clear();
    Ballot = prop.num_prop;
    Ballot.append("&&");
    Ballot.append(prop.sub_prop).append("&&");
    Ballot.append(prop.desc_prop).append("&&");
    Ballot.append(prop.tipo_voto).append("&&");

    // genera una ordenamiento aleatorio de la boleta que se le envia al usuario
    switch(prop.tipo_voto.toInt()){
    case 1:
        {
            QString num="12";
            while (num.size()>0){
                pointer =qrand() % num.size();
                switch(num.at(pointer).toAscii()){
                    case '1':
                        Ballot.append("A favor").append("&&");
                        break;
                    case '2':
                        Ballot.append("En contra").append("&&");
                        break;
                }
               num.remove(pointer,1);
            }
            Ballot.append(prop.misc);
        }
        break;
    case 2:
        {
            QString num="12";
            while (num.size()>0){
                pointer =qrand() % num.size();
                switch(num.at(pointer).toAscii()){
                    case '1':
                        Ballot.append(prop.opc1_desc).append("&&");
                        break;
                    case '2':
                        Ballot.append(prop.opc2_desc).append("&&");
                        break;
                }
                num.remove(pointer,1);
            }
            Ballot.append(prop.misc);
        }
        break;
    case 3:
        {
            QString num="123";
            while (num.size()>0){
                pointer =qrand() % num.size();
                switch(num.at(pointer).toAscii()){
                    case '1':
                        Ballot.append(prop.opc1_desc).append("&&");
                        break;
                    case '2':
                        Ballot.append(prop.opc2_desc).append("&&");
                        break;
                    case '3':
                        Ballot.append(prop.opc3_desc).append("&&");
                        break;
                }
                num.remove(pointer,1);
            }
            Ballot.append(prop.misc);
        }
        break;
    case 4:
        {
            QString num="1234";
            while (num.size()>0){
                pointer =qrand() % num.size();
                switch(num.at(pointer).toAscii()){
                    case '1':
                        Ballot.append(prop.opc1_desc).append("&&");
                        break;
                    case '2':
                        Ballot.append(prop.opc2_desc).append("&&");
                        break;
                    case '3':
                        Ballot.append(prop.opc3_desc).append("&&");
                        break;
                    case '4':
                        Ballot.append(prop.opc4_desc).append("&&");
                        break;
                }
                num.remove(pointer,1);
            }
            Ballot.append(prop.misc);
        }
        break;
    case 5:
        {
            QString num="12345";
            while (num.size()>0){
                pointer =qrand() % num.size();
                switch(num.at(pointer).toAscii()){
                    case '1':
                        Ballot.append(prop.opc1_desc).append("&&");
                        break;
                    case '2':
                        Ballot.append(prop.opc2_desc).append("&&");
                        break;
                    case '3':
                        Ballot.append(prop.opc3_desc).append("&&");
                        break;
                    case '4':
                        Ballot.append(prop.opc4_desc).append("&&");
                        break;
                    case '5':
                        Ballot.append(prop.opc5_desc).append("&&");
                        break;
                }
                num.remove(pointer,1);
            }
            Ballot.append(prop.misc);
        }
        break;
    }

}

void ballot_generator::create_key(){
    QString namePubKey="eb";
    QString namePriKey="db";

    InvertibleRSAFunction params;
    params.GenerateRandomWithKeySize( rng, 1024);
    RSA::PrivateKey pri( params );

    RSA::PublicKey pub(pri);

    SavePublicKey(namePubKey,pub);
    SavePrivateKey(namePriKey,pri);
}

void ballot_generator::SavePrivateKey(QString &filename, const PrivateKey &key){
    ByteQueue queue;
    key.Save(queue);

    Save(filename, queue);
}

void ballot_generator::SavePublicKey(QString &filename, const PublicKey &key){
    ByteQueue queue;
    key.Save(queue);

    Save(filename, queue);
}

void ballot_generator::Save(QString &filename, const BufferedTransformation &bt){
    filename.prepend("/");
    filename.prepend(obj_Dir.absolutePath());

    FileSink file(filename.toAscii().data());

    bt.CopyTo(file);
    file.MessageEnd();
}

void ballot_generator::LoadPublicKey(QString filename, PublicKey &key){
    // http://www.cryptopp.com/docs/ref/class_byte_queue.html
    ByteQueue queue;
    Load(filename, queue);
    key.Load(queue);
}

void ballot_generator::LoadPrivateKey(QString filename, PrivateKey &key){
    // http://www.cryptopp.com/docs/ref/class_byte_queue.html
    ByteQueue queue;
    Load(filename, queue);

    key.Load(queue);
}

void ballot_generator::Load(QString filename, BufferedTransformation& bt)
{
    // http://www.cryptopp.com/docs/ref/class_file_source.html
    filename.prepend("/");
    filename.prepend(obj_Dir.absolutePath());

    FileSource file(filename.toAscii().data(), true /*pumpAll*/);
    file.TransferTo(bt);

    bt.MessageEnd();

}

/*bool ballot_generator::PVID_verification(QString PVID){
// Crea la conexion al Padron Propuesta (PdnP)
    QSqlDatabase PVID_List = QSqlDatabase::database("PVIDconex",false);
    QDir PVID_Dir(obj_Dir);
    PVID_Dir.cdUp();
    PVID_List.setDatabaseName(PVID_Dir.absoluteFilePath("PVID_List.db"));

    if(!PVID_List.open()){
            return false;
        }

// Crea el modelo
    QSqlTableModel *PVIDmodel = new QSqlTableModel(0,PVID_List);
    PVIDmodel->setTable("list");
    PVIDmodel->setEditStrategy(QSqlTableModel::OnManualSubmit);
    PVIDmodel->setSort(1,Qt::AscendingOrder);
    PVIDmodel->select();

    for(int i = 0; i < PVIDmodel->rowCount();i++){
        string decodedKey;
        HexDecoder deco(new StringSink(decodedKey));
        deco.Put((byte*)PVID.toStdString().data(),PVID.size());
        deco.MessageEnd();

        HMAC<SHA> t((const byte*)PVIDmodel->record(i).value("pin").toString().toAscii().data(),PVIDmodel->record(i).value("pin").toString().size());

        if(t.VerifyDigest((byte*)decodedKey.data(),
                           (const byte*)PVIDmodel->record(i).value("RFTag").toString().toAscii().data(),
                           PVIDmodel->record(i).value("RFTag").toString().size())){
            if(PVIDmodel->record(i).value("vote_saved").toString()==QString("no"))
                return true;
            else
                return false;
            }
    }
    return false;
}*/
